# filepath: d:\CompanyProject\Gitee\DBCompareX\backend\app\models\task_data_sync\db_connectors.py
"""
数据库连接器模块，负责处理与不同类型数据库的连接和查询
"""
from typing import List, Dict, Any, Tuple
import logging
import time
import pymysql
import pyodbc

from ...models.db_info_setting import DBInfoSettingModel

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class DBConnector:
    """数据库连接器，提供统一的数据库查询接口"""
    
    # 重试配置
    MAX_RETRIES = 3
    RETRY_DELAY = 2
    
    @classmethod
    async def execute_query_with_retry(cls, db_info: Dict[str, Any], sql: str, params: List = None) -> List[Tuple]:
        """
        带重试机制的数据库查询
        
        Args:
            db_info: 数据库配置信息
            sql: SQL查询语句
            params: 查询参数
            
        Returns:
            查询结果
        """
        retries = 0
        last_error = None
        
        while retries < cls.MAX_RETRIES:
            try:
                return await cls.execute_query(db_info, sql, params)
            except Exception as e:
                last_error = e
                retries += 1
                
                if retries < cls.MAX_RETRIES:
                    logger.warning(f"查询失败，准备第 {retries} 次重试: {str(e)}")
                    # 等待一段时间后重试
                    time.sleep(cls.RETRY_DELAY)
                else:
                    logger.error(f"查询失败，已达到最大重试次数 {cls.MAX_RETRIES}: {str(e)}")
                    raise last_error
    
    @classmethod
    async def execute_query(cls, db_info: Dict[str, Any], sql: str, params: List = None) -> List[Tuple]:
        """
        执行数据库查询
        
        Args:
            db_info: 数据库配置信息
            sql: SQL查询语句
            params: 查询参数
            
        Returns:
            查询结果
        """
        db_type = db_info["db_type"].lower()
        
        if db_type == "mysql":
            return await cls.execute_mysql_query(db_info, sql, params)
        elif db_type == "sqlserver":
            return await cls.execute_sqlserver_query(db_info, sql, params)
        elif db_type == "fabric_sqlserver":
            return await cls.execute_fabric_sqlserver_query(db_info, sql, params)
        else:
            raise ValueError(f"不支持的数据库类型: {db_type}")
            
    @classmethod
    async def execute_fabric_sqlserver_query(cls, db_info: Dict[str, Any], sql: str, params: List = None) -> List[Tuple]:
        """
        执行Microsoft Fabric Data Warehouse查询
        
        Args:
            db_info: 数据库配置信息
            sql: SQL查询语句
            params: 查询参数
            
        Returns:
            查询结果
        """
        try:
            # 构建Fabric专用连接字符串
            conn_str = (
                f"DRIVER={{ODBC Driver 18 for SQL Server}};"  # 尽可能使用18版驱动
                f"SERVER={db_info['host']};"  # Fabric端点通常不需要端口
                f"DATABASE={db_info['database_name']};"
                f"Uid={db_info['username']};"
                f"Pwd={db_info['password']};"
                f"Encrypt=yes;"
                f"TrustServerCertificate=yes;"
                f"Authentication=ActiveDirectoryServicePrincipal;"
                f"ssl=true;"
                f"Connection Timeout=300;"
            )
            
            # 添加额外的连接参数
            if db_info.get("connection_string_ext"):
                conn_str += f"{db_info['connection_string_ext']}"
            logger.info(f"连接 Microsoft Fabric Data Warehouse: {db_info['host']}, 数据库: {db_info['database_name']}")
            connection = pyodbc.connect(conn_str)
            try:
                with connection.cursor() as cursor:
                    cursor.execute(sql, params or ())
                    return cursor.fetchall()
            finally:
                connection.close()
        except Exception as e:
            logger.error(f"Microsoft Fabric Data Warehouse查询失败: {str(e)}", exc_info=True)
            raise ValueError(f"Microsoft Fabric Data Warehouse查询失败: {str(e)}")
    
    @classmethod
    async def execute_mysql_query(cls, db_info: Dict[str, Any], sql: str, params: List = None) -> List[Tuple]:
        """
        执行MySQL查询
        
        Args:
            db_info: 数据库配置信息
            sql: SQL查询语句
            params: 查询参数
            
        Returns:
            查询结果
        """
        try:
            connection = pymysql.connect(
                host=db_info["host"],
                port=int(db_info["port"]),
                user=db_info["username"],
                password=db_info["password"],
                database=db_info["database_name"],
                connect_timeout=30,  # 连接超时时间
                read_timeout=30,     # 读取超时时间
                write_timeout=30     # 写入超时时间
            )
            
            try:
                with connection.cursor() as cursor:
                    cursor.execute(sql, params or ())
                    return cursor.fetchall()
            finally:
                connection.close()
        except Exception as e:
            logger.error(f"MySQL查询失败: {str(e)}", exc_info=True)
            raise ValueError(f"MySQL查询失败: {str(e)}")
    
    @classmethod
    async def execute_sqlserver_query(cls, db_info: Dict[str, Any], sql: str, params: List = None) -> List[Tuple]:
        """
        执行SQL Server查询
        
        Args:
            db_info: 数据库配置信息
            sql: SQL查询语句
            params: 查询参数
            
        Returns:
            查询结果
        """
        try:
            # 构建连接字符串
            conn_str = (
                f"DRIVER={{ODBC Driver 18 for SQL Server}};"
                f"SERVER={db_info['host']},{db_info['port']};"
                f"DATABASE={db_info['database_name']};"
                f"UID={db_info['username']};"
                f"PWD={db_info['password']};"
                f"Connection Timeout=300;"
                f"TrustServerCertificate=yes;"
                f"Query Timeout=300"
            )
            
            # 添加额外的连接参数
            if db_info.get("connection_string_ext"):
                conn_str += f";{db_info['connection_string_ext']}"
                
            logger.info(f"连接 SQL Server: {db_info['host']}:{db_info['port']}, 数据库: {db_info['database_name']}")
            
            connection = pyodbc.connect(conn_str)
            
            try:
                with connection.cursor() as cursor:
                    cursor.execute(sql, params or ())
                    return cursor.fetchall()
            finally:
                connection.close()
        except Exception as e:
            logger.error(f"SQL Server查询失败: {str(e)}", exc_info=True)
            raise ValueError(f"SQL Server查询失败: {str(e)}")